/**
   @file     iot_shell.c
   @brief    远程shell调试通道 (字符编码UTF8)
   @author   zxx (zxx@comtika.com)
   @date     2025-03-31
   @version  v1.0.0
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>
#include <time.h>
#include <ctype.h>
#include <pthread.h>
#include "iot_shell.h"


#define BUFFER_SIZE 4096
#define READ_END  0
#define WRITE_END 1
#define CMD_END_MARKER  "##CMD_EXECUTION_END##"
#define CMD_TIMEOUT_MS  5000   // 命令执行超时时间(5s)

// 全局变量
volatile pid_t shell_pid = -1;  // 子进程PID
int parent_to_child[2];         // 主进程向子进程传递命令的管道
int child_to_parent[2];         // 子进程向主进程传递结果的管道

// 线程锁
pthread_mutex_t _lock = PTHREAD_MUTEX_INITIALIZER;

//////////////////////////////////////////////////////////

/**
   @brief 超时检测
*/
static uint8_t check_timeout(struct timespec start_time) {
    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);
    long ms = (now.tv_sec - start_time.tv_sec) * 1000;
    ms += (now.tv_nsec - start_time.tv_nsec) / 1000000;
    return ms > CMD_TIMEOUT_MS;
}

/**
   @brief shell是否已启动
   @return 1-true，0-false
*/
uint8_t is_shell_open()
{
    if (shell_pid <= 0) return 0;
    else return 1;
}

/**
   @brief 启动shell子进程 
   @return 1-成功，0-失败
*/
uint8_t shell_open()
{
    pthread_mutex_lock(&_lock);
    if (is_shell_open()) {
        pthread_mutex_unlock(&_lock);
        return 1; // 子进程已启动
    }

    // 创建管道
    if (pipe(parent_to_child) < 0 || pipe(child_to_parent) < 0) {
        pthread_mutex_unlock(&_lock);
        perror("shell: 管道创建失败");
        return 0;
    }
    
    // 创建子进程
    shell_pid = fork();
    if (shell_pid < 0) {
        pthread_mutex_unlock(&_lock);
        perror("shell: fork失败");
        return 0;
    }
    
    // -------------- 子进程 --------------
    if (shell_pid == 0) { 
        // 关闭不需要的管道端
        close(parent_to_child[WRITE_END]);
        close(child_to_parent[READ_END]);
        
        // 重定向标准输入输出
        dup2(parent_to_child[READ_END], STDIN_FILENO);
        dup2(child_to_parent[WRITE_END], STDOUT_FILENO);
        dup2(child_to_parent[WRITE_END], STDERR_FILENO);
        
        // 启动shell终端
        execl("/bin/sh", "sh", NULL);
        
        // 如果启动成功，下面的代码不会被执行
        perror("shell: execl失败");
        exit(EXIT_FAILURE);
    }
    // -------------- 父进程 --------------
    else {
        // 关闭不需要的管道端
        close(parent_to_child[READ_END]);
        close(child_to_parent[WRITE_END]);
        
        // 设置子进程到父进程的管道为非阻塞模式
        fcntl(child_to_parent[READ_END], F_SETFL, O_NONBLOCK);

        pthread_mutex_unlock(&_lock);

        printf("shell: pid=%d, name=[sh]\n", shell_pid);
        return 1;
    }
}

/**
   @brief 关闭shell子进程
*/
void shell_close()
{
    pthread_mutex_lock(&_lock);
    if (!is_shell_open()) {
        pthread_mutex_unlock(&_lock);
        return;
    }
    
    // 关闭管道
    if (parent_to_child[WRITE_END] >= 0) {
        close(parent_to_child[WRITE_END]);
        parent_to_child[WRITE_END] = -1;
    }
    if (child_to_parent[READ_END] >= 0) {
        close(child_to_parent[READ_END]);
        child_to_parent[READ_END] = -1;
    }
    
    // 发送终止信号给子进程
    pid_t pid = shell_pid;
    kill(pid, SIGTERM);

    shell_pid = -1; // 重置PID后解锁
    pthread_mutex_unlock(&_lock);
    
    // 等待子进程结束
    int status;
    pid_t result = waitpid(pid, &status, WNOHANG);
    
    // 如果子进程没有立即结束，给它一些时间，等待最多3秒
    if (result == 0) {
        int wait_count = 0;
        while (wait_count < 30) {
            usleep(100000); // 等待0.1秒
            result = waitpid(pid, &status, WNOHANG);
            if (result != 0) {
                break;
            }
            wait_count++;
        }
        
        // 如果子进程仍未结束，强制终止
        if (result == 0) {
            kill(pid, SIGKILL);
            waitpid(pid, &status, 0);
        }
    }
}

/**
   @brief 检查是否进程启动命令
*/
static uint8_t is_access_command(const char *cmd) {
    // 检查是否以 ./ 或 / 或 ~ 开头（常见 Shell 启动方式）
    if (strncmp(cmd, "./", 2) == 0 || cmd[0] == '/' || cmd[0] == '~') {
        char cmd_head[50];

        size_t head_len = strcspn(cmd, " "); // 计算第一个空格前命令的长度
        strncpy(cmd_head, cmd, 50);
        cmd_head[head_len] = '\0';
        printf("is_access_command? = %s, %d\n", cmd_head, access(cmd_head, X_OK));
        if (access(cmd_head, X_OK) == 0) {
            return 1;
        }
    }
    return 0;
}

/**
   @brief 检查命令是否为危险命令或不支持的命令
   @note  1. 交互式命令（如vim、nano等）
          2. 危险命令（如rm -rf、mkfs等）
          3. 需要root权限的命令（如sudo、su、passwd等）
*/
static int8_t check_cmmand(const char *cmd, char *output, size_t output_size)
{
    if (cmd == NULL || strlen(cmd) <= 1) {
        snprintf(output, output_size, "无效的参数");
        return -4;
    }

    if (is_access_command(cmd)) {
        snprintf(output, output_size, "检测到进程启动命令: %s, 已被拦截", cmd);
        return -1;
    }

    // 检查是否为交互式命令
    const char *interactive_cmds[] = {"vim", "vi", "nano", "top", "htop", "less", "more"};
    for (int i = 0; i < sizeof(interactive_cmds) / sizeof(interactive_cmds[0]); i++) {
        if (strstr(cmd, interactive_cmds[i]) != NULL) {
            snprintf(output, output_size, "不支持交互式命令: %s", interactive_cmds[i]);
            return -1;
        }
    }
    
    // 检查是否为危险命令
    const char *dangerous_cmds[] = {"rm -rf", "mkfs", "dd", "fdisk", "fsck", "format", "> /dev/", "mv /", "chmod -R 777 /", "/etc/passwd", "kill -9"};
    for (int i = 0; i < sizeof(dangerous_cmds) / sizeof(dangerous_cmds[0]); i++) {
        if (strstr(cmd, dangerous_cmds[i]) != NULL) {
            snprintf(output, output_size, "警告: 检测到可能有风险的命令: %s, 已被拦截", dangerous_cmds[i]);
            return -2;
        }
    }

    // 拦截需要root权限的命令
    const char *root_cmds[] = {"sudo", "su", "passwd", "chown", "adduser", "deluser", "useradd", "userdel", "groupmod", "usermod", 
		"mount", "umount", "ulimit", "systemctl", "service", "iptables", "firewall-cmd", "visudo", "poweroff", "shutdown", "halt", "init"};
    for (int i = 0; i < sizeof(root_cmds) / sizeof(root_cmds[0]); i++) {
        if (strstr(cmd, root_cmds[i]) != NULL) {
            snprintf(output, output_size, "警告: 检测到需要root权限的命令: %s, 已被拦截", root_cmds[i]);
            return -3;
        }
    }
    
    return 0; // 命令通过检查
}

/**
   @brief 命令格式化，防止Shell注入风险

   @param cmd       输入命令
   @param out_cmd   输出格式化后的命令
   @param out_size  输出缓冲区最大长度
   @return          返回命令长度
*/
uint32_t format_command(const char *cmd, char *out_cmd, size_t out_size)
{
    int j = 0;              // 新字符串索引
    int prev_is_space = 1;  // 标记前一个字符是否已被处理为空格

    strncpy(out_cmd, cmd, out_size);
    for (int i = 0; out_cmd[i] != '\0'; i++) {
        if (out_cmd[i] == '\n') {
            out_cmd[i] = '\0';
            break;
        }
        if (isspace((unsigned char)out_cmd[i])) {
            if (!prev_is_space) {
                out_cmd[j++] = ' ';
                prev_is_space = 1;
            }
        } else {
            out_cmd[j++] = out_cmd[i];
            prev_is_space = 0;
        }
    }
    out_cmd[strcspn(out_cmd, ";")] = '\0';
    return strlen(out_cmd);
}

/**
   @brief   执行shell命令并获取结果
   @return  错误码: 0-成功，<0失败
*/
int8_t shell_exec_cmd(const char *cmd, char *output, size_t output_size) {
    char buffer[BUFFER_SIZE];

    if (cmd == NULL || output == NULL || output_size <= 0) {
        return -4;
    }
    if ( !is_shell_open() ) {
        snprintf(output, output_size, "远程终端已退出!");
        return -1;
    }

    output[0] = '\0';
    size_t cmd_len = format_command(cmd, buffer, sizeof(buffer));
    printf("CMD: %s\n", buffer);
    
    // 检查命令安全性
    int err_check = check_cmmand(buffer, output, output_size);
    if (err_check != 0) {
        return err_check;
    }

    // 构造带有结束标记的命令
    snprintf(buffer+cmd_len, sizeof(buffer)-cmd_len, "; echo '%s'\n", CMD_END_MARKER);
    
    // 发送命令到子进程
    write(parent_to_child[WRITE_END], buffer, strlen(buffer));
    usleep(10000);  // 等待读取结果
    
    
    int bytes_read;
    uint8_t find_end = 0;
    size_t output_len = 0;
    char output_buffer[BUFFER_SIZE * 2] = {0}; // 用于存储累积的输出

    struct timespec start_time;
    clock_gettime(CLOCK_MONOTONIC, &start_time);
    
    while (!find_end) {
        if (check_timeout(start_time)) {
            if (output_len + 30 < output_size) {
                strcat(output, "执行超时(5s),终端关闭");
            }
            shell_close();
            return -1;
        }
        
        memset(buffer, 0, BUFFER_SIZE);
        bytes_read = read(child_to_parent[READ_END], buffer, BUFFER_SIZE - 1);
        
        if (bytes_read > 0) {
            buffer[bytes_read] = '\0';

            // 查找是否包含结束标记
            char *end_pos = strstr(buffer, CMD_END_MARKER);
            if (end_pos != NULL) {
                find_end = 1;
                bytes_read = (end_pos - buffer);
                *end_pos = '\0';
            }
            
            // 将新读取的内容追加到累积缓冲区
            if (output_len + bytes_read < sizeof(output_buffer)) {
                memcpy(output_buffer + output_len, buffer, bytes_read);
                output_len += bytes_read;
                output_buffer[output_len] = '\0';
            } else {
                find_end = 1; // 缓冲区溢出
                bytes_read = sizeof(output_buffer) - output_len -1;
                memcpy(output_buffer + output_len, buffer, bytes_read);
                output_len += bytes_read;
                output_buffer[output_len] = '\0';

                do {  // 清空缓冲区
                    bytes_read = read(child_to_parent[READ_END], buffer, sizeof(buffer));
                    usleep(10000);
                    if (check_timeout(start_time)) {
                        shell_close();
                        return -1;
                    }
                } while (bytes_read > 0);
            }
            
            if (find_end) {
                // 复制结果到输出缓冲区
                strncpy(output, output_buffer, output_size - 1);
                output[output_size - 1] = '\0';
                break;
            }
        } else if (bytes_read == -1 && errno == EAGAIN) {
            // 没有更多数据可读，等待一下
            usleep(10000);
        } else {
            snprintf(output, output_size, "读取命令输出错误");
            shell_close();
            return -2;
        }
    }
    
    return 0;
}


/************ 测试用例 ************
int main_test() {
    char cmd[BUFFER_SIZE];
    char output_buff[SHELL_BUFF_SIZE];
    
    // 初始化shell子进程
    if (!shell_open()) {
        fprintf(stderr, "Shell初始化失败\n");
        exit(EXIT_FAILURE);
    }
    
    printf("自定义Shell启动 (输入'exit'退出)\n");
    
    while (1) {
        printf("> ");
        fflush(stdout);
        
        // 读取用户输入
        if (fgets(cmd, BUFFER_SIZE, stdin) == NULL) {
            break;
        }
        
        // 执行命令
        shell_exec_cmd(cmd, output_buff, SHELL_BUFF_SIZE);
        printf("%s", output_buff);

        // 是否退出
        if (!is_shell_open()) {
            break;
        }
    }
    
    printf("Shell已退出\n");
    
    return 0;
}
**********************************/
